"""

Hybrid A* path planning for onsite

author: tpr@sjtu

"""

import heapq
import numpy as np
import math
from math import sqrt, cos, sin, tan, pi

try:
    import hybrid_a_star.reeds_shepp_path_planning as rs
except:
    raise


XY_GRID_RESOLUTION = 2.0  # [m]
YAW_GRID_RESOLUTION = np.deg2rad(15.0)  # [rad]
MOTION_RESOLUTION = 0.1  # [m] path interporate resolution
N_STEER = 20  # number of steer command

WB = 5.6  #轴距
W = 5.5  # width of car
LF = 7.2  # distance from rear to vehicle front end
LB = 1.6  # distance from rear to vehicle back end
# MAX_STEER = np.deg2rad(30)  # [rad] maximum steering angle

# SB_COST = 100.0  # switch back penalty cost
BACK_COST = 5.0  # backward penalty cost
SB_COST = np.inf  # 换倒挡惩罚无限大
BACK_COST = np.inf  # 倒车惩罚无限大
STEER_CHANGE_COST = 5.0  # steer angle change penalty cost
STEER_COST = 1.0  # steer angle change penalty cost
H_COST = 5.0  # Heuristic cost



class Node:

    def __init__(self, xind, yind, yawind, direction,
                 xlist, ylist, yawlist, directions,
                 steer=0.0, pind=None, cost=None):
        self.xind = xind
        self.yind = yind
        self.yawind = yawind
        self.direction = direction
        self.xlist = xlist
        self.ylist = ylist
        self.yawlist = yawlist
        self.directions = directions
        self.steer = steer
        self.pind = pind
        self.cost = cost


class Path:

    def __init__(self, xlist, ylist, yawlist, directionlist, cost):
        self.xlist = xlist
        self.ylist = ylist
        self.yawlist = yawlist
        self.directionlist = directionlist
        self.cost = cost

class Config:

    def __init__(self, observation, xyreso, yawreso):
        min_x_m = observation['test_setting']['x_min']
        min_y_m = observation['test_setting']['y_min']
        max_x_m = observation['test_setting']['x_max']
        max_y_m = observation['test_setting']['y_max']

        self.minx = round(min_x_m / xyreso) # 网格的最小和最大x、y索引
        self.miny = round(min_y_m / xyreso) # 将observation字典中的最小和最大x、y值除以xyreso(分辨率)参数并四舍五入计算得到
        self.maxx = round(max_x_m / xyreso)
        self.maxy = round(max_y_m / xyreso)

        self.xw = round(self.maxx - self.minx) # 网格的宽度和高度
        self.yw = round(self.maxy - self.miny)

        self.minyaw = round(- math.pi / yawreso) - 1
        self.maxyaw = round(math.pi / yawreso)
        self.yaww = round(self.maxyaw - self.minyaw)


def calc_motion_inputs(): # 生成车辆可能的转向角和方向,范围从-MAX_STEER到MAX_STEER，生成的值的数量为N_STEER

    for steer in np.concatenate((np.linspace(-MAX_STEER, MAX_STEER, N_STEER),[0.0])): # 将额外的转向角0.0附加到范围上
        for d in [1]: # 只有单个方向[1]，表示前进
            yield [steer, d]

def check_car_collision(xlist, ylist, yawlist, collision_lookup,observation,image_ndarray_new): # 检查给定路径是否会导致碰撞
    local_x_range = observation['hdmaps_info']['image_mask'].bitmap_info['bitmap_mask_PNG']['UTM_info']['local_x_range']
    local_y_range = observation['hdmaps_info']['image_mask'].bitmap_info['bitmap_mask_PNG']['UTM_info']['local_y_range']
    for x, y, yaw in zip(xlist, ylist, yawlist):    # 通过减去局部范围的最小x和y值来调整x和y坐标
        if collision_lookup.collision_detection(x-local_x_range[0],
                                                y-local_y_range[0],
                                                yaw,
                                                image_ndarray_new):
            return True
    return False

def pi_2_pi(angle):
    return (angle + pi) % (2 * pi) - pi

def move(x, y, yaw, distance, steer, L=WB):
    x += distance * cos(yaw)     # 根据转向角和移动距离计算车辆的新位置和方向
    y += distance * sin(yaw)
    yaw += pi_2_pi(distance * tan(steer) / L)  # distance/2

    return x, y, yaw

def get_neighbors(current, config,collision_lookup,observation,image_ndarray_new,MOTION_RESOLUTION):    # 从给定节点生成所有可能的下一个节点

    for steer, d in calc_motion_inputs():
        node = calc_next_node(current, steer, d, config, collision_lookup,observation,image_ndarray_new,MOTION_RESOLUTION)
        if node and verify_index(node, config):
            yield node


def calc_next_node(current, steer, direction, config, collision_lookup,observation,image_ndarray_new,MOTION_RESOLUTION):    # 根据当前节点、转向角和方向计算下一个节点

    x, y, yaw = current.xlist[-1], current.ylist[-1], current.yawlist[-1]   # 先从当前节点中提取最后的x、y和偏航值

    arc_l = XY_GRID_RESOLUTION * 1.5    # 弧长arc_l表示在转向角度方向上移动的距离
    xlist, ylist, yawlist = [], [], []
    for dist in np.arange(0, arc_l, MOTION_RESOLUTION):
        x, y, yaw = move(x, y, yaw, MOTION_RESOLUTION * direction, steer)   # 从当前位置朝转向角的方向移动来计算下一个节点的新x、y和偏航值
        xlist.append(x)
        ylist.append(y)
        yawlist.append(yaw)
    if check_car_collision(xlist, ylist, yawlist, collision_lookup,observation,image_ndarray_new):
        return None

    d = direction
    xind = round(x / XY_GRID_RESOLUTION)
    yind = round(y / XY_GRID_RESOLUTION)
    yawind = round(yaw / YAW_GRID_RESOLUTION)

    addedcost = 0.0

    if d != current.direction:  # 换挡惩罚, 100
        addedcost += SB_COST

    addedcost += STEER_COST * abs(steer)    # 偏航值惩罚
    addedcost += STEER_CHANGE_COST * abs(current.steer - steer)    # 偏航值改变惩罚

    cost = current.cost + addedcost + arc_l

    node = Node(xind, yind, yawind, d, xlist,
                ylist, yawlist, [d],
                pind=calc_index(current, config),
                cost=cost, steer=steer)

    return node


def is_same_grid(n1, n2):
    if n1.xind == n2.xind and n1.yind == n2.yind and n1.yawind == n2.yawind:
        return True
    return False

# import matplotlib.pyplot as plt
# import os
# test_plot_path = os.path.join(os.path.dirname(__file__), 'test_plot_path.png')
# plt.plot(path.x, path.y, label='test_plot_path')
# plt.legend()
# plt.savefig(test_plot_path, dpi=300)
# print(f'路径已保存到{test_plot_path}')

def analytic_expantion(current, goal, c, collision_lookup, observation,bo_alp_in,image_ndarray_new,MOTION_RESOLUTION):

    sx = current.xlist[-1]  # 获取当前节点的最后一个x坐标
    sy = current.ylist[-1]  # 获取当前节点的最后一个y坐标
    syaw = current.yawlist[-1]  # 获取当前节点的最后一个偏航角

    gx = goal.xlist[-1]  # 获取目标节点的最后一个x坐标
    gy = goal.ylist[-1]  # 获取目标节点的最后一个y坐标
    gyaw = goal.yawlist[-1]  # 获取目标节点的最后一个偏航角

    max_curvature = math.tan(MAX_STEER) / WB  # 计算最大曲率，用于Reeds-Shepp路径规划
    paths = rs.calc_paths(sx, sy, syaw, gx, gy, gyaw,
                          max_curvature, step_size=MOTION_RESOLUTION)  # 使用Reeds-Shepp算法计算可能的路径

    if not paths:  # 如果没有找到任何路径
        return None

    best_path, best = None, None
    for path in paths:  # 遍历所有可能的路径
        # 首先检查路径是否包含倒车段
        has_backward = False
        for l in path.lengths:
            if l < 0:  # 如果长度小于0，表示这段是倒车
                has_backward = True
                break
        
        # 只考虑不包含倒车的路径
        if not has_backward and not check_car_collision(path.x, path.y, path.yaw, collision_lookup,observation,image_ndarray_new):
            cost = calc_rs_path_cost(path)  # 计算路径代价
            if not best or best > cost:  # 如果这是第一条路径或代价更小
                best = cost  # 更新最小代价
                best_path = path  # 更新最佳路径
    
    return best_path  # 返回最佳路径


def update_node_with_analystic_expantion(current, goal,    # 定义函数，接收当前节点、目标节点等参数
                                         c, collision_lookup, observation,bo_alp_in,image_ndarray_new,MOTION_RESOLUTION):
    # bo_alp_in参数将被忽略，始终不允许倒车
    apath = analytic_expantion(current, goal, c, collision_lookup, observation,False,image_ndarray_new,MOTION_RESOLUTION)    # 尝试使用Reeds-Shepp路径从当前节点到目标节点
    
    if apath:    # 如果找到了可行路径
        fx = apath.x[1:]    # 获取路径的x坐标（去掉第一个点，因为与当前节点重复）
        fy = apath.y[1:]    # 获取路径的y坐标
        fyaw = apath.yaw[1:]    # 获取路径的偏航角

        fcost = current.cost + calc_rs_path_cost(apath)    # 计算总代价：当前节点代价加上Reeds-Shepp路径代价
        fpind = calc_index(current, c)    # 计算当前节点的索引值

        # 确保所有方向都是前进
        fd = [True] * len(apath.directions[1:])

        fsteer = 0.0    # 设置转向角为0
        fpath = Node(current.xind, current.yind, current.yawind,    # 创建新的路径节点
                     True, fx, fy, fyaw, fd,
                     cost=fcost, pind=fpind, steer=fsteer)
        return True, fpath    # 返回成功标志和路径节点

    return False, None    # 如果没有找到可行路径，返回失败标志和None


def calc_rs_path_cost(rspath):

    # 首先检查是否有倒车路径，如果有直接返回无穷大代价
    for l in rspath.lengths:
        if l < 0:  # 如果是倒车，直接返回无穷大代价
            return float('inf')
    
    cost = 0.0
    
    # 正常计算前进路径的代价
    for l in rspath.lengths:
        cost += l  # 只有前进路径，所以直接累加长度
    
    # steer penalty
    for ctype in rspath.ctypes:
        if ctype != "S":  # curve
            cost += STEER_COST * abs(MAX_STEER)

    # steer change penalty
    # calc steer profile
    nctypes = len(rspath.ctypes)
    ulist = [0.0] * nctypes
    for i in range(nctypes):
        if rspath.ctypes[i] == "R":
            ulist[i] = - MAX_STEER
        elif rspath.ctypes[i] == "L":
            ulist[i] = MAX_STEER

    for i in range(len(rspath.ctypes) - 1):
        cost += STEER_CHANGE_COST * abs(ulist[i + 1] - ulist[i])

    return cost


def hybrid_a_star_planning(start, goal, collision_lookup, observation, xyreso, yawreso,bo_alp_in,image_ndarray_new,iter_number,MOTION_RESOLUTION,max_steer=30):
    """
    start
    goal
    xyreso: grid resolution [m]
    yawreso: yaw angle resolution [rad]
    """
    # bo_alp_in参数将被忽略，路径规划始终不允许倒车
    global MAX_STEER
    MAX_STEER = np.deg2rad(max_steer)
    print("Start Hybrid A* planning!")
    start[2], goal[2] = rs.pi_2_pi(start[2]), rs.pi_2_pi(goal[2])

    config = Config(observation, xyreso, yawreso)

    # 确保起始节点和目标节点的方向都是前进
    nstart = Node(round(start[0] / xyreso), round(start[1] / xyreso), round(start[2] / yawreso),
                  True, [start[0]], [start[1]], [start[2]], [True], cost=0)
    ngoal = Node(round(goal[0] / xyreso), round(goal[1] / xyreso), round(goal[2] / yawreso),
                 True, [goal[0]], [goal[1]], [goal[2]], [True])

    # 初始化了用于跟踪尚未访问的节点和已访问的节点的开放列表和关闭列表
    openList, closedList = {}, {}

    pq = []
    openList[calc_index(nstart, config)] = nstart
    heapq.heappush(pq, (calc_cost(nstart,goal),calc_index(nstart, config)))
    iter_num = 0
    while True:
        iter_num +=1
        if not openList:
            print("Error: Cannot find path, No open set")
            return [], [], []

        cost, c_id = heapq.heappop(pq)
        if c_id in openList:
            current = openList.pop(c_id)
            closedList[c_id] = current
        else:
            continue

        # 尝试使用Reeds-Shepp路径直接连接到目标（强制不允许倒车）
        isupdated, fpath = update_node_with_analystic_expantion(
            current, ngoal, config, collision_lookup, observation,False,image_ndarray_new,MOTION_RESOLUTION)

        if isupdated:
            print("use rs successfully!")
            break

        # 生成邻居节点（只考虑前进方向）
        for neighbor in get_neighbors(current, config, collision_lookup,observation,image_ndarray_new,MOTION_RESOLUTION):
            neighbor_index = calc_index(neighbor, config)
            if neighbor_index in closedList:
                continue
            if neighbor not in openList \
                    or openList[neighbor_index].cost > neighbor.cost:
                heapq.heappush(pq, (calc_cost(neighbor,goal),neighbor_index))
                openList[neighbor_index] = neighbor
        
        if iter_num>iter_number:
            break
    print(f"iter_num = {iter_num}")
    path = get_final_path(closedList, fpath, nstart, config)
    return path


def calc_cost(n,goal):
    cost = sqrt((n.xlist[0] - goal[0]) **2 + (n.ylist[0] - goal[1]) **2)
    return n.cost + H_COST * cost


def get_final_path(closed, ngoal, nstart, config):
    if ngoal is None:
        return None
    else:
        rx, ry, ryaw = list(reversed(ngoal.xlist)), list(
            reversed(ngoal.ylist)), list(reversed(ngoal.yawlist))
        direction = list(reversed(ngoal.directions))
        # print(len(rx),len(direction))
        nid = ngoal.pind
        finalcost = ngoal.cost

        while nid:
            n = closed[nid]
            # print(n.directions)
            rx.extend(list(reversed(n.xlist)))
            ry.extend(list(reversed(n.ylist)))
            ryaw.extend(list(reversed(n.yawlist)))
            # 确保所有方向都是前进
            direction.extend(list(reversed([True] * len(n.xlist))))
            nid = n.pind
        # print(len(rx),len(direction))
        rx = list(reversed(rx))
        ry = list(reversed(ry))
        ryaw = list(reversed(ryaw))
        direction = list(reversed(direction))
        # print(direction)
        # 确保所有方向都是前进
        for i in range(len(direction)):
            direction[i] = True

        path = Path(rx, ry, ryaw, direction, finalcost)

    return path


def verify_index(node, c):
    xind, yind = node.xind, node.yind
    if xind >= c.minx and xind <= c.maxx and yind >= c.miny \
            and yind <= c.maxy:
        return True

    return False


def calc_index(node, c):
    ind = (node.yind - c.miny) * c.xw + (node.xind - c.minx)

    if ind <= 0:
        print("Error(calc_index):", ind)

    return ind
